#include "opticalflow.h"


OpticalFlow::OpticalFlow::OpticalFlow() {}
OpticalFlow::~OpticalFlow()
{
    if( u )
        cvReleaseImage(&u);
    if( v )
        cvReleaseImage(&v);
    if( ex )
        cvReleaseImage(&ex);
    if( ey )
        cvReleaseImage(&ey);
    if( et)
        cvReleaseImage(&et);
    if( t1 )
        cvReleaseImage(&t1);
    if( t2 )
        cvReleaseImage(&t2);
    if( uL )
        cvReleaseImage(&uL);
    if( vL )
        cvReleaseImage(&vL);
}
OpticalFlow::OpticalFlow(const OpticalFlow &) {}
OpticalFlow& OpticalFlow::operator=(const OpticalFlow &) {return (*this);}

inline void OpticalFlow::estimateOpticalFlow(IplImage *frame)
{
    cvConvertImage(frame,t2);
    unsigned int iter(0), j(0), i(0) , width(frame->width), height(frame->height);
    double error(max_error+1.0);

    char  * pt1, *pt2;
    double * pex, * pey, * pet, * pu, * pv , * puL, * pvL, den, nom, uTMP, vTMP;

    //calculs des Ex, Ey, Et....
    pt1 =          t1->imageData + width + 1;
    pt2 =          t2->imageData + width + 1;
    pex = (double*)(ex->imageData + width + 1 );
    pey = (double*)(ey->imageData + width + 1 );
    pet = (double*)(et->imageData + width + 1 );
    for(j=1 ; j < height-1 ; j++)
    {
        for(i=1 ; i < width-1 ; i++)
        {
            *(pex++) = 0.25*((double)(*(pt1+width)) - (double)(*(pt1))   + (double)(*(pt1+width+1)) - (double)(*(pt1+1))
                           + (double)(*(pt2+width)) - (double)(*(pt2))   + (double)(*(pt2+width+1)) - (double)(*(pt2+1)));
            *(pey++) = 0.25*((double)(*(pt1+1))     - (double)(*(pt1))   + (double)(*(pt1+width+1)) - (double)(*(pt1+width))
                           + (double)(*(pt2+1))     - (double)(*(pt2))   + (double)(*(pt2+width+1)) - (double)(*(pt2+width)));
            *(pet++) = 0.25*((double)(*(pt2))       + (double)(*(pt2+1)) + (double)(*(pt2+width))   + (double)(*(pt2+width+1))
                           - (double)(*(pt1))       - (double)(*(pt1+1)) - (double)(*(pt1+width))   - (double)(*(pt1+width+1)));

            pt1++;
            pt2++;
        }
        pt1+=2;
        pt2+=2;
        pex+=2;
        pey+=2;
        pet+=2;
    }


    cvSetZero(u);
    cvSetZero(v);

    double a[] = { 1.0/12.0, 1.0/6.0 , 1.0/12.0,
                   1.0/6.0 , 0.0    ,  1.0/6.0 ,
                   1.0/12.0, 1.0/6.0 , 1.0/12.0 };

    CvMat Laplace=cvMat(3, 3, CV_64FC1, a);


    for(; (iter < max_iter) /*&& (error > max_error )*/ ; iter++)
    {
        cvFilter2D(u,uL,&Laplace);
        cvFilter2D(v,vL,&Laplace);

        error = 0.0;
        pt1 =          t1->imageData + width + 1;
        pt2 =          t2->imageData + width + 1;
        pex = (double*)(ex->imageData + width + 1 );
        pey = (double*)(ey->imageData + width + 1 );
        pet = (double*)(et->imageData + width + 1 );
        pu  = (double*)( u->imageData + width + 1 );
        pv  = (double*)( v->imageData + width + 1 );
        puL = (double*)(uL->imageData + width + 1 );
        pvL = (double*)(vL->imageData + width + 1 );
        for(j=1 ; j < height-1 ; j++)
        {
            for(i=1 ; i < width-1 ; i++)
            {
                uTMP = *puL;
                vTMP = *pvL;

                den = (*pex)*uTMP + (*pey)*vTMP + (*pet);
                nom = (alpha2 + (*pex)*(*pex) + (*pey)*(*pey))*9.0;

                *(pu) = uTMP - (*pex)*den/nom;
                *(pv) = vTMP - (*pey)*den/nom;


                error+= (*pu-uTMP)*(*pu-uTMP) + (*pv-vTMP)*(*pv-vTMP);

                pt1++;
                pt2++;
                pex++;
                pey++;
                pet++;
                pu++;
                pv++;
                puL++;
                pvL++;
            }
            pt1+=2;
            pt2+=2;
            pex+=2;
            pey+=2;
            pet+=2;
            pu+=2;
            pv+=2;
            puL+=2;
            pvL+=2;
        }

        cout<<iter+1<<" : "<<error<<endl;
    }

    cvCopy(t2,t1);
}

void OpticalFlow::estimateResult()
{
        int i,j;

        double * pu  = (double*)( u->imageData   + u->width   + 1 );
        double * pv  = (double*)( v->imageData   + v->width   + 1 );
        char  * pR  =         ( res->imageData + res->width + 1 );
        double tmp;
        for(j=1 ; j < res->height-1 ; j++)
        {
            for(i=1 ; i < res->width-1 ; i++)
            {
                tmp = (*pu)*(*pu) + (*pv)*(*pv);
                *(pR++) = ((tmp>threshold)?200:0);
                pu++;
                pv++;
            }
            pu+=2;
            pv+=2;
            pR+=2;
        }

}

IplImage * OpticalFlow::init  (IplImage *frame)
{
    u   = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);
    cvSetZero(u);
    v   = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);
    cvSetZero(v);
    uL   = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);
    vL   = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);

    res = cvCreateImage(cvSize(frame->width,frame->height), 8, 1);
    cvSetZero(res);

    ex  = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);
    ey  = cvCreateImage(cvSize(frame->width,frame->height), 64, 1);
    et  = cvCreateImage(cvSize(frame->width,frame->height), 64 , 1);
    t1  = cvCreateImage(cvSize(frame->width,frame->height), 8, 1);
    t2  = cvCreateImage(cvSize(frame->width,frame->height), 8, 1);

    max_iter  = 100 ;
    max_error = 10e-10;
    alpha2    = 0.0001;
    threshold = 15.0;

    return res;
}

IplImage * OpticalFlow::update(IplImage *frame)
{
    estimateOpticalFlow(frame);
    estimateResult();
    /*cvErode( res,t2);
    cvDilate( t2,res);*/
    return res;
}
